
#include <openssl/objects.h>
#include <apr-1/apr_strings.h>
#include "../iface/debug.h"
#include "../utils/Hash.h"
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "SSHClient.h"

SSHClient::~SSHClient()
{
   if (our_pubkey) delete our_pubkey;
   if (our_prvkey) delete our_prvkey;
   if (our_pass) { memset((void *)our_pass, 0, strlen(our_pass)); }
   if (tchannel)
   {
      libssh2_channel_free(tchannel);
      tchannel = NULL;
   }
   if (tsession)
   {
      libssh2_session_disconnect(tsession, "Finished session");
      libssh2_session_free(tsession);
   }
}

void *SSHClient::KeyboardInteractiveHandler(const char *name, int name_len,
                         const char *instruction, int instruction_len,
                         int num_prompts,
                         const LIBSSH2_USERAUTH_KBDINT_PROMPT *prompts,
                         LIBSSH2_USERAUTH_KBDINT_RESPONSE *responses,
                         void **abstract)
{
   SSHClient *kthis = (SSHClient *)(*abstract);
   (void)name;
   (void)name_len;
   (void)instruction;
   (void)instruction_len;
   if (num_prompts == 1)
   {
      responses[0].text = strdup(kthis->our_pass);
      responses[0].length = strlen(kthis->our_pass);
   }
   (void)prompts;
   (void)abstract;
}

int SSHClient::PublicKeyHandler(LIBSSH2_SESSION *session,
                                  unsigned char **sig,
                                  size_t *sig_len,
                                  const unsigned char *data,
                                  size_t data_len,
                                  void **abstract)
{
   SSHClient *kthis = (SSHClient *)(*abstract);
   char *hex_sha1 = Hash::getSHA1Hash((const char *)data, data_len);
   int sha1_hash_bytes = strlen(hex_sha1);
   char *sha1_hash = BaseConvert::convertHexadecimalToBytes(hex_sha1, &sha1_hash_bytes);
   free(hex_sha1);
   unsigned int signature_size = RSA_size(kthis->our_prvkey->prv);
   unsigned char *signature_blob = (unsigned char *)malloc(signature_size);
   if (signature_blob == NULL) { free(sha1_hash); return -1; }
   int retval = RSA_sign(NID_sha1, (const unsigned char *)sha1_hash, sha1_hash_bytes, signature_blob, &signature_size, kthis->our_prvkey->prv);
   free(sha1_hash);
   if (!retval) { free(signature_blob); return -1; }
   *sig = signature_blob;
   *sig_len = signature_size;
   return 0;
}

const char *SSHClient::DoHandshakeAndGetHostFingerprintInHex()
{
   // SHA-1 is 160bit, 20bytes
   // MD-5 is 128bit, 16bytes
   const char *sha1hashtmp = DoHandshakeAndGetHostFingerprint();
   if (sha1hashtmp == NULL)
   {
      errprint("handshake failed");
      return NULL;
   }
   char *sha1hexhash = BaseConvert::convertBytesToHexadecimal(sha1hashtmp, 20);
   return sha1hexhash;
}

const char *SSHClient::DoHandshakeAndGetHostFingerprint()
{
   tsession = libssh2_session_init_ex(NULL, NULL, NULL, this);
   if (tsession == NULL) return NULL;
   apr_status_t status;
   int berk_socket_des;
   status = apr_os_sock_get(&berk_socket_des, tsocket);
   CheckAPRError(status);
   if (libssh2_session_handshake(tsession, berk_socket_des))
   {
      errprint("failed to achieve handshake with server");
      return NULL;
   }
   return libssh2_hostkey_hash(tsession, LIBSSH2_HOSTKEY_HASH_SHA1);
}

void SSHClient::SetAuthenticatorKeys(KeyDB *sshkeys, const char *keyname, const char *keypass)
{
   our_pubkey = (RSAKeyPair *)sshkeys->GetPublicKeyByName(keyname);
   our_prvkey = (RSAKeyPair *)sshkeys->GetPrivateKeyByName(keyname, keypass);
   char *tmp_keyname = (char *)apr_palloc(mp, sizeof(char)*(strlen(keyname)+1));
   strcpy(tmp_keyname, keyname);
   our_keyname = tmp_keyname;
}

void SSHClient::SetAuthenticatorPassword(const char *authpass)
{
   char *tmp_pass = (char *)apr_palloc(mp, sizeof(char)*(strlen(authpass)+1));
   strcpy(tmp_pass, authpass);
   our_pass = (const char *)tmp_pass;
}

#define USE_AUTHENTICATION_TYPE_PASSWORD 1
#define USE_AUTHENTICATION_TYPE_KEY      2

int SSHClient::ForcePasswordAuthentication(const char *username)
{
   return SendUsernameAndDoAuthentication(username, USE_AUTHENTICATION_TYPE_PASSWORD);
}

int SSHClient::ForceKeyAuthentication(const char *username)
{
   return SendUsernameAndDoAuthentication(username, USE_AUTHENTICATION_TYPE_KEY);
}

int SSHClient::SendUsernameAndDoAuthentication(const char *username)
{
   return SendUsernameAndDoAuthentication(username, USE_AUTHENTICATION_TYPE_PASSWORD|USE_AUTHENTICATION_TYPE_KEY);
}

int SSHClient::SendUsernameAndDoAuthentication(const char *username, int authflags)
{
   if (username == NULL)
   {
      errprint("no username given");
      return 1;
   }
   char *tmp_username = (char *)apr_palloc(mp, sizeof(char)*(strlen(username)+1));
   strcpy(tmp_username, username);
   our_username = tmp_username;
   char *userauthlist = libssh2_userauth_list(tsession, username, strlen(username));
   if ((strstr(userauthlist, "password") != NULL) && (our_pass != NULL) && (authflags & USE_AUTHENTICATION_TYPE_PASSWORD))
   {
      return libssh2_userauth_password(tsession, username, our_pass);
   }
   if ((strstr(userauthlist, "keyboard-interactive") != NULL) && (our_pass != NULL) && (authflags & USE_AUTHENTICATION_TYPE_PASSWORD))
   {
      return libssh2_userauth_keyboard_interactive(tsession, username, GetKeyboardInteractiveHandlerFunction());
   }
   if ((strstr(userauthlist, "publickey") != NULL) && (our_prvkey != NULL) && (authflags & USE_AUTHENTICATION_TYPE_KEY))
   {
      void *abstract = (void *)this;
      int pubkeybloblen = 0;
      char *pubkeyblob = our_pubkey->ExportOpenSSHPublicKeyBlob(&pubkeybloblen);
      int retval = libssh2_userauth_publickey(tsession, username, (const unsigned char *)pubkeyblob, pubkeybloblen, GetPublicKeyHandlerFunction(), &abstract);
      free(pubkeyblob);
      return retval;
   }
   return 1;
}

int SSHClient::OpenTunnelChannel(const char *remote_host, int remote_port, int local_port)
{
   tchannel = libssh2_channel_direct_tcpip_ex(tsession, remote_host, remote_port, "127.0.0.1", local_port);
   if (tchannel == NULL)
   {
      errprint("could not establish tunnel on port %d (ssh -L %d:%s:%d)", local_port, local_port, remote_host, remote_port);
      return -1;
   }
   libssh2_session_set_blocking(tsession, 0);
   return 0;
}

int SSHClient::CopyFile(const char *local_file, const char *remote_file)
{
}

void SSHClient::TrimRightWhitespace(char *thisline)
{
   // Do right trim of whitespace (<= ' ')
   for (int i = strlen(thisline); ((i--) && (thisline[i] <= ' ')); ) thisline[i] = '\0';
}

#define RUNCMDBUFFERSIZE 4096

char *SSHClient::RunCommand(const char *command)
{
   std::vector<char> *output_buffer = new std::vector<char>();
   tchannel = libssh2_channel_open_session(tsession);
   if (libssh2_channel_exec(tchannel, command))
   {
      errprint("could not execute '%s'", command);
      char *errmsg;
      int errlen;
      int err = libssh2_session_last_error(tsession, &errmsg, &errlen, 0);
      errprint("ssh error: %s", errmsg);
      delete output_buffer;
      return NULL;
   }
   char *readbuf = (char *)apr_palloc(mp, sizeof(char)*(4096));
   int rc = 1;
   while (rc > 0)
   {
      rc = libssh2_channel_read(tchannel, readbuf, RUNCMDBUFFERSIZE);
      if (rc > 0)
         for (int i = 0; i < rc; i++) 
            output_buffer->push_back(readbuf[i]);
      if (rc == LIBSSH2_ERROR_EAGAIN)
      {
         apr_sleep(1000);
         rc = 1;
      }
   }
   output_buffer->push_back('\0');
   char *outputbuf = (char *)apr_palloc(mp, sizeof(char)*(strlen(&(output_buffer->at(0)))+1));
   strcpy(outputbuf, &(output_buffer->at(0)));
   delete output_buffer;
   TrimRightWhitespace(outputbuf);
   libssh2_channel_flush(tchannel);
   libssh2_channel_close(tchannel);
   libssh2_channel_free(tchannel);
   tchannel = NULL;
   return outputbuf;
}

int SSHClient::PutDataIntoSecureRemoteFile(const char *dir, const char *file, const char *stringbuf, int stringlen)
{
   char *direcho = StringUtils::mpprintf(mp, "echo %s/%s", dir, file);
   char *outputbuf = RunCommand(direcho);
   if (outputbuf == NULL)
   {
      errprint("could not obtain remote value for '%s/%s'", dir, file);
      return -1;
   } else
   {
      errprint("copying across file to %s", outputbuf);
   }
   tchannel = libssh2_scp_send64(tsession, outputbuf, 0700, stringlen, time(NULL), time(NULL));
   if (tchannel == NULL)
   {
      errprint("could not open scp channel, copying to %s", outputbuf);
      char *errmsg;
      int errlen;
      int err = libssh2_session_last_error(tsession, &errmsg, &errlen, 0);
      errprint("ssh error: %s", errmsg);
      return -1;
   }
   ssize_t written = libssh2_channel_write_ex(tchannel, 0, stringbuf, stringlen);
   libssh2_channel_flush(tchannel);
   libssh2_channel_close(tchannel);
   libssh2_channel_free(tchannel);
   tchannel = NULL;
   return 1;
}

/*
 *    tchannel = libssh2_scp_send64(tsession, brokerkeyfilepath, 0700, strlen(newkeyline), time(NULL), time(NULL));
   if (tchannel == NULL)
   {
      errprint("could not open scp channel, copying to %s", brokerkeyfilepath);
      char *errmsg;
      int errlen;
      int err = libssh2_session_last_error(tsession, &errmsg, &errlen, 0);
      errprint("ssh error: %s", errmsg);
      return -1;
   }
   ssize_t written = libssh2_channel_write_ex(tchannel, 0, newkeyline, strlen(newkeyline));
   libssh2_channel_flush(tchannel);
   libssh2_channel_close(tchannel);
*/

int SSHClient::DeleteThisKey(const char *oldkeyname)
{
   char *outputbuf = RunCommand("echo $HOME");
   if (outputbuf == NULL)
   {
      errprint("could not obtain value for $HOME");
      return -1;
   }

   char *oldkeytag = StringUtils::mpprintf(mp, "ssi-amrmmhd.epcc.ed.ac.uk_%s", oldkeyname);
   char *authkeysfilepath = StringUtils::mpprintf(mp, "%s/.ssh/authorized_keys", outputbuf);

   char *buildcmd_4 = StringUtils::mpprintf(mp, "grep \"%s\" %s", oldkeytag, authkeysfilepath);
   char *outputbuf_4 = RunCommand(buildcmd_4);
   if ((outputbuf_4 == NULL) || (strlen(outputbuf_4) < strlen(oldkeytag)))
   {
      if (outputbuf_4 == NULL)
      {
         errprint("could not check keys");
         return -1;
      } else
      {
         if (strlen(outputbuf_4) == 0) errprint("could not find key '%s' present for deletion from remote host", oldkeyname);
         if (strlen(outputbuf_4) > 0) errprint("possibly could not find key '%s', response '%s'", oldkeyname, outputbuf_4);
         return 1;
      }
   }

   char *buildcmd_2 = StringUtils::mpprintf(mp, "sh -c '\
      touch %s; chmod 700 %s; set -C;\
      if echo \"$$\" > /tmp/authorized_keys_lock; then\
         grep -v \"%s\" %s | tee %s;\
         rm /tmp/authorized_keys_lock;\
      fi'", authkeysfilepath, authkeysfilepath, oldkeytag, authkeysfilepath, authkeysfilepath);
   char *outputbuf_2 = RunCommand(buildcmd_2);
   if (outputbuf_2 == NULL)
   {
      errprint("update to authorized_keys failed");
      return -1;
   }

   char *buildcmd_3 = StringUtils::mpprintf(mp, "grep \"%s\" %s", oldkeytag, authkeysfilepath);
   char *outputbuf_3 = RunCommand(buildcmd_3);
   if ((outputbuf_3 == NULL) || (strlen(outputbuf_3) > strlen(oldkeytag)))
   {
      if (outputbuf_3 == NULL)
         errprint("could not check key update");
      else
      {
         if (strlen(outputbuf_3) > strlen(oldkeytag)) errprint("key update failed");
      }
      return -1;
   }

   return 0;
}

int SSHClient::AddThisKey(const char *newkeyname, RSAKeyPair *newpubkey)
{
   char *outputbuf = RunCommand("echo $HOME");
   if (outputbuf == NULL)
   {
      errprint("could not obtain value for $HOME");
      return -1;
   }
   
   char *newkeytag = StringUtils::mpprintf(mp, "ssi-amrmmhd.epcc.ed.ac.uk_%s", newkeyname);
   char *newkeyline = NULL;
   char *tmp_newkeyline = newpubkey->ExportOpenSSHPublicKey(newkeytag);
   if (tmp_newkeyline == NULL)
   {
      errprint("could not get public key '%s'", newkeytag);
      return -1;
   } else
   {
      newkeyline = apr_pstrdup(mp, tmp_newkeyline);
      free(tmp_newkeyline);
   }
   char *authkeysfilepath = StringUtils::mpprintf(mp, "%s/.ssh/authorized_keys", outputbuf);

   char *buildcmd_5 = StringUtils::mpprintf(mp, "grep \"%s\" %s", newkeytag, authkeysfilepath);
   char *outputbuf_5 = RunCommand(buildcmd_5);
   if ((outputbuf_5 == NULL) || (strlen(outputbuf_5) > strlen(newkeytag)))
   {
      if (outputbuf_5 == NULL)
         errprint("could not check keys");
      else
      {
         if (strlen(outputbuf_5) > strlen(newkeytag))
            errprint("key update already happened");
         return 0;
      }
      return -1;
   }
   
   char *buildcmd_2 = StringUtils::mpprintf(mp, "sh -c '\
      touch %s; chmod 700 %s; set -C;\
      if echo \"$$\" > /tmp/authorized_keys_lock; then\
         echo -n \"%s\n\" >> %s;\
         rm /tmp/authorized_keys_lock;\
      fi'", authkeysfilepath, authkeysfilepath, newkeyline, authkeysfilepath);
   char *outputbuf_2 = RunCommand(buildcmd_2);
   if (outputbuf_2 == NULL)
   {
      errprint("update to authorized_keys failed");
      return -1;
   }

   char *buildcmd_3 = StringUtils::mpprintf(mp, "grep \"%s\" %s", newkeytag, authkeysfilepath);
   char *outputbuf_3 = RunCommand(buildcmd_3);
   if ((outputbuf_3 == NULL) || (strlen(outputbuf_3) < strlen(newkeytag)))
   {
      if (outputbuf_3 == NULL)
         errprint("could not check key update");
      else
      {
         if (strlen(outputbuf_3) == 0)
         {
            errprint("key update failed");
            return 1;
         }
         if (strlen(outputbuf_3) > 0) errprint("confusing response '%s'", outputbuf_3);
      }
      return -1;
   }
   return 0;
}

LIBSSH2_CHANNEL *SSHClient::GetSSHClientChannel()
{
   return tchannel;
}

//bash -c 'set -C; if echo "$$" > /tmp/authorized_keys_lock; then grep -v "myname" /root/.ssh/authorized_keys > /root/.ssh/authorized_keys; echo "newkey" >> /root/.ssh/authorized_keys; rm /tmp/authorized_keys_lock; fi'
